#include "cross_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
namespace optiling
{
  const uint32_t BLOCK_SIZE = 64;
  static ge::graphStatus TilingFunc(gert::TilingContext *context)
  {
    CrossTilingData tiling;
    uint32_t sizeofdatatype;
    uint32_t totalLengthAligned;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto socVersion = ascendcPlatform.GetSocVersion();
    uint64_t ub_size;
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
    auto aivNum = ascendcPlatform.GetCoreNumAiv(); // vector core num  1
    auto x1Shape = context->GetInputTensor(0)->GetShape().GetOriginShape();
    auto x2Shape = context->GetInputTensor(1)->GetShape().GetOriginShape();
    int64_t dim_attr = *context->GetAttrs()->GetInt(0);
    // bool isSameShape = true;
    // if (x1Shape.GetDimNum() == x2Shape.GetDimNum())
    // {
    //   for (int i = x1Shape.GetDimNum() - 1; i >= 0; i--)
    //   {
    //     if (x1Shape.GetDim(i) != x2Shape.GetDim(i))
    //     {
    //       isSameShape = false;
    //       break;
    //     }
    //   }
    // }
    // else
    // {
    //   isSameShape = false;
    // }
    // if(isSameShape){
    uint32_t x1iValue = 1;
    uint32_t iValue = 1;
    uint32_t x1jValue = 1;
    uint32_t jValue = 1;
    uint32_t x2iValue = 1;
    uint32_t x2jValue = 1;
    uint32_t dim = -1;
    uint32_t totalLength = 1;
    uint32_t dimNum = x1Shape.GetDimNum();
    uint32_t *yshape = new uint32_t[dimNum];
    uint32_t *x1shape = new uint32_t[dimNum];
    uint32_t *x2shape = new uint32_t[dimNum];
    if (dim_attr == 65530 || dim_attr >= dimNum || dim_attr < 0)
    { // 寻找第一个维度均为3的维度
      for (int i = 0; i < dimNum; i++)
      {
        uint32_t _dim = x1Shape.GetDim(i);
        uint32_t min_dim = x2Shape.GetDim(i);
        if (min_dim > _dim)
        {
          _dim = min_dim;
          min_dim = x1Shape.GetDim(i);
        }
        totalLength *= _dim;
        if (dim == -1 && (_dim != 3 || min_dim != 3))
        {
          iValue *= _dim;
          x1iValue *= x1Shape.GetDim(i);
          x2iValue *= x2Shape.GetDim(i);
          yshape[i] = _dim;
          x1shape[i] = x1Shape.GetDim(i);
          x2shape[i] = x2Shape.GetDim(i);
        }
        else if (dim != -1)
        {
          jValue *= _dim;
          x1jValue *= x1Shape.GetDim(i);
          x2jValue *= x2Shape.GetDim(i);
          yshape[i] = _dim;
          x1shape[i] = x1Shape.GetDim(i);
          x2shape[i] = x2Shape.GetDim(i);
        }
        else
        {
          dim = i;
          yshape[i] = 3;
          x1shape[i] = 3;
          x2shape[i] = 3;
        }
      }
    }
    else
    { // 按照指定的维度进行叉乘
      for (int i = 0; i < dimNum; i++)
      {
        dim = dim_attr;
        uint32_t _dim = x1Shape.GetDim(i);
        if (x2Shape.GetDim(i) > _dim)
        {
          _dim = x2Shape.GetDim(i);
        }
        totalLength *= _dim;
        if (i < dim_attr)
        {
          iValue *= _dim;
          x1iValue *= x1Shape.GetDim(i);
          x2iValue *= x2Shape.GetDim(i);
          yshape[i] = _dim;
          x1shape[i] = x1Shape.GetDim(i);
          x2shape[i] = x2Shape.GetDim(i);
        }
        else if (i > dim_attr)
        {
          jValue *= _dim;
          x1jValue *= x1Shape.GetDim(i);
          x2jValue *= x2Shape.GetDim(i);
          yshape[i] = _dim;
          x1shape[i] = x1Shape.GetDim(i);
          x2shape[i] = x2Shape.GetDim(i);
        }
        else
        {
          yshape[i] = 3;
          x1shape[i] = 3;
          x2shape[i] = 3;
        }
      }
    }
    uint32_t x1Length = context->GetInputTensor(0)->GetShapeSize();
    uint32_t x2Length = context->GetInputTensor(1)->GetShapeSize();
    auto dt = context->GetInputTensor(0)->GetDataType();
    uint32_t typeKey = 0;
    uint32_t data_num = 2;
    if (dt == 1)
    {                     // half
      sizeofdatatype = 2; // 输出需要是int8类型的
      tiling.set_typeKey(1);
    }
    else if (dt == 0)
    { // float
      sizeofdatatype = 4;
      tiling.set_typeKey(0);
    }
    else if (dt == 2)
    { // int8
      sizeofdatatype = 1;
      tiling.set_typeKey(2);
    }
    else if (dt == 3)
    { // int32
      sizeofdatatype = 4;
      tiling.set_typeKey(3);
    }
    uint32_t ALIGN_NUM = BLOCK_SIZE / sizeofdatatype;
    // 核内拆分，策略是尽可能的填满ub_size，最后一包单独处理，
    uint32_t ub_block_num =
        ((ub_size) / BLOCK_SIZE / data_num) * 13 / 20; // ub_block在Ascend C中不能全部被用来作为输入输出，给了13/20系数。
    // printf("ub_size:%lu\t\t%u*************\n", ub_size, ub_block_num);
    // uint32_t ub_block_num = 5;  //为测试方便，验证代码流程
    uint32_t tile_num;

    context->SetBlockDim(1);
    if (totalLength % ALIGN_NUM != 0)
    { // 不对齐，先32位对齐
      totalLengthAligned =
          ((totalLength + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
    }
    else
    {
      totalLengthAligned = totalLength;
    }
    uint32_t tileLength = 0;
    uint32_t lasttileLength = 0;
    tile_num = totalLengthAligned / ALIGN_NUM / ub_block_num; // 简单获取tile_num的大小
    if ((totalLengthAligned / ALIGN_NUM) % ub_block_num == 0 ||
        tile_num == 0)
    { // 对齐之后，可以均分到各个tile，或不足一个tile
      if (tile_num == 0)
      {
        tile_num = 1;
      }
      if (totalLengthAligned < ub_block_num * ALIGN_NUM)
      {
        tileLength = totalLengthAligned;
        lasttileLength = tileLength;
      }
      else
      {
        tileLength = ub_block_num * ALIGN_NUM;
        lasttileLength = tileLength;
      }
    }
    else
    { // 满足字节对齐，核内不能均分
      tile_num = tile_num + 1;
      tileLength = ub_block_num * ALIGN_NUM;
      lasttileLength = totalLengthAligned - (tile_num - 1) * tileLength;
    }
    printf("%u\t\t%u******%u*******\n", tile_num, tileLength, lasttileLength);
    // printf("%u\t\t%u*************\n", dim, dimNum);
    printf("%u\t\t%u*************\n", iValue, jValue);
    // printf("%u\t\t%u*************\n", x1iValue, x1jValue);
    // printf("%u\t\t%u*************\n", x2iValue, x2jValue);
    // for(int i=0;i<dimNum;i++){
    //   printf("%u\t\t%u******%u*******\n", x1shape[i], x2shape[i], yshape[i]);
    // }
    tiling.set_totalLength(totalLength);
    tiling.set_dim(dim);
    tiling.set_yshape(yshape);
    tiling.set_x1shape(x1shape);
    tiling.set_x2shape(x2shape);
    tiling.set_dimNum(dimNum);
    tiling.set_x1Length(x1Length);
    tiling.set_x2Length(x2Length);
    tiling.set_tileNum(tile_num);
    tiling.set_tileLength(tileLength);
    tiling.set_lasttileLength(lasttileLength);
    tiling.set_x1iValue(x1iValue);
    tiling.set_x1jValue(x1jValue);
    tiling.set_x2iValue(x2iValue);
    tiling.set_x2jValue(x2jValue);
    tiling.set_jValue(jValue);
    tiling.set_iValue(iValue);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
                        context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    return ge::GRAPH_SUCCESS;
    // }
  }
}

namespace ge
{
  static ge::graphStatus InferShape(gert::InferShapeContext *context)
  {
    const gert::Shape *x1_shape = context->GetInputShape(0);
    gert::Shape *y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
  }
}

namespace ops
{
  class Cross : public OpDef
  {
  public:
    explicit Cross(const char *name) : OpDef(name)
    {
      this->Input("x1")
          .ParamType(REQUIRED)
          .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT8, ge::DT_INT32})
          .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
          .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
      this->Input("x2")
          .ParamType(REQUIRED)
          .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT8, ge::DT_INT32})
          .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
          .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
      this->Output("y")
          .ParamType(REQUIRED)
          .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT8, ge::DT_INT32})
          .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
          .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
      this->Attr("dim").AttrType(OPTIONAL).Int(65530);

      this->SetInferShape(ge::InferShape);

      this->AICore()
          .SetTiling(optiling::TilingFunc);
      this->AICore().AddConfig("ascend310b");
    }
  };

  OP_ADD(Cross);
}
